from PIL import Image, ImageDraw
import fire


def get_data(file_name):
    with open(file_name, "rb") as f:
        return f.read()


def write_byte(co, bt):
    co = list(co)
    co[0] = co[0] & 0b11111110 | ((bt >> 0) & 1)
    co[1] = co[1] & 0b11111110 | ((bt >> 1) & 1)
    co[2] = co[2] & 0b11111110 | ((bt >> 2) & 1)
    co[3] = co[3] & 0b11111110 | ((bt >> 3) & 1)
    return tuple(co)


def read_byte(co):
    r, g, b, a = co[0] & 1, co[1] & 1, co[2] & 1, co[3] & 1
    return r | (g << 1) | (b << 2) | (a << 3)


def expand_bytes(bs):
    r = []
    for item in bs:
        r.append(item & 0b00001111)
        r.append(item >> 4)
    return r


def shrink_bytes(bs):
    r = [0 for _ in range(len(bs)//2)]
    for i in range(0, len(bs)//2):
        r[i] = bs[2*i] | (bs[2*i+1] << 4)
    return r


def encode(input_file, output_file, data):
    im = Image.open(input_file).convert("RGBA")
    width, height = im.size
    if width*height // 2 < len(data) + 8:
        return "图片太小，无法隐藏足够的信息"
    content = []
    sz = len(data)
    for i in range(16):
        content.append((sz >> (i * 4)) & 0xf)
    for it in expand_bytes(data):
        content.append(it)
    out = Image.new('RGBA', (width, height))
    draw = ImageDraw.Draw(out)
    cnt = 0
    for x in range(width):
        for y in range(height):
            if cnt < len(content):
                co = im.getpixel((x, y))
                draw.point((x, y), fill=write_byte(co, content[cnt]))
            else:
                co = im.getpixel((x, y))
                draw.point((x, y), fill=co)
            cnt += 1
    out.save(output_file, 'png')


def decode(input_file, src_file):
    im = Image.open(input_file).convert("RGBA")
    width, height = im.size
    sz = 0
    for i in range(16):
        x, y = i / height, i % height
        co = im.getpixel((x, y))
        bt = read_byte(co)
        sz = sz | (bt << (i*4))
    data = []
    for i in range(16, 16 + sz*2):
        x, y = i / height, i % height
        co = im.getpixel((x, y))
        bt = read_byte(co)
        data.append(bt)
    data = shrink_bytes(data)
    with open(src_file, "wb") as fp:
        fp.write(bytearray(data))

def read_command(source_file = "", output_file = "", hide_file = ""):
    if not source_file:
        print("please input the source file name")
        return
    if not output_file:
        print("please input the output file name")
        return
    if hide_file:
        file_data = get_data(hide_file)
        if r := encode(source_file, output_file, file_data):
            print(r)
    else:
        decode(source_file, output_file)

if __name__ == "__main__":
    fire.Fire(read_command)